import torch.utils.data as data
from os import listdir
from os.path import join
from PIL import Image,ImageFilter
import numpy as np
import  matplotlib.pyplot as plt
from torchvision.transforms import Compose,CenterCrop,ToTensor,Scale


def is_image_file(filename):
    if (filename.endswith(extension) for extension in [".png",".jpg","tiff"]):
        return True
    else:
        return False

def load_img(filepath):
    img = Image.open(filepath).convert('YCbCr')
    y, _, _ = img.split()
    return y
    #return img

class My_datasets(data.Dataset):

    def __init__(self,img_dir,input_transform =None,target_transform =None):
        super(My_datasets,self).__init__()
        self.image_filenames = [join(img_dir,x) for x in listdir(img_dir) if is_image_file(x) ]
        self.input_transform = input_transform   
        self.target_transform = target_transform 

    def __getitem__(self, index):
        input = load_img(self.image_filenames[index])
        target = input.copy()   
        # print(type(input))
        if self.input_transform:
            input = input.filter(ImageFilter.GaussianBlur(3))   

            input = np.array(input)
            H,W = input.shape
            mean,sigma = 0,0.02
            noise = np.random.normal(mean,sigma,input.shape)
            noise = noise*255
            input = input+noise

            input = Image.fromarray(input.astype('uint8'))

            input = self.input_transform(input)

        if self.target_transform:
            target = self.target_transform(target)
        return input, target

    def __len__(self):
        return len(self.image_filenames)


if __name__ == '__main__':
    p = My_datasets('./BSDS300/images/train')

    print(p.__getitem__(1)[0]) 
    plt.subplot(121)
    plt.imshow(p.__getitem__(1)[0])
    plt.subplot(122)
    plt.imshow(p.__getitem__(1)[1])
    plt.show()



